{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from modeling_llama import LlamaForCausalLM\n",
    "from transformers import AutoTokenizer\n",
    "from decoding import clip_input, infer_input_ids\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_path = './CodeLlama-13b/'\n",
    "torch.nn.Linear.reset_parameters = lambda x: None\n",
    "model = LlamaForCausalLM.from_pretrained(model_path, torch_dtype=torch.bfloat16)\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device='cuda:2'\n",
    "model = model.to(device).eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_attn_skip_layer_id_set, _mlp_skip_layer_id_set =model.get_skip_layers()\n",
    "print(_attn_skip_layer_id_set, _mlp_skip_layer_id_set)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "seed=42\n",
    "torch.manual_seed(seed)\n",
    "np.random.seed(seed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from human_eval.data import write_jsonl\n",
    "from tqdm import tqdm\n",
    "def filter_code(completion: str) -> str:\n",
    "    # The program tends to overwrite, we only take the first function\n",
    "    completion = completion.lstrip(\"\\n\")\n",
    "    return completion.split(\"\\n\\n\")[0]\n",
    "\n",
    "def fix_indents(text: str) -> str:\n",
    "    return text.replace(\"\\t\", \"    \")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "from human_eval.data import read_problems\n",
    "n_shot = 0\n",
    "task_name = 'humaneval'\n",
    "prompt_shots = ''\n",
    "if task_name == 'xsum':\n",
    "    data = load_dataset('xsum', split='test').shuffle(seed=seed).select(range(1000))\n",
    "    shots = load_dataset('xsum',split='train').shuffle(seed=seed).select(range(n_shot))\n",
    "    prompt_keys=['document','summary']\n",
    "elif task_name == 'cnndm':\n",
    "    data = load_dataset('cnn_dailymail', name='3.0.0', split='test') .shuffle(seed=seed).select(range(1000))\n",
    "    shots = load_dataset('cnn_dailymail', name='3.0.0', split='train').shuffle(seed=seed).select(range(n_shot))\n",
    "    prompt_keys=['article','highlights']\n",
    "elif task_name == 'humaneval':\n",
    "    data = read_problems() # adopt hf load humaneval have problems\n",
    "    \n",
    "for i in range(n_shot):\n",
    "    prompt = 'Article: ' + shots[i][prompt_keys[0]] + '\\nSummary: ' + shots[i][prompt_keys[1]].replace('\\n', '') + '\\n'\n",
    "    prompt_shots += prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(prompt_shots)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "main_metrics = {'time_base':[], 'time_sss':[], 'time_sss_autoth1':[], 'time_sss_autoth2':[],\n",
    "                'token_time_base':[], 'token_time_sss':[], 'token_time_sss_autoth1':[], 'token_time_sss_autoth2':[],\n",
    "                'matchness_sss':[],'num_drafted_tokens_sss':[],\n",
    "                'matchness_sss_autoth1':[],'num_drafted_tokens_sss_autoth1':[],\n",
    "                'matchness_sss_autoth2':[],'num_drafted_tokens_sss_autoth2':[],\n",
    "                }\n",
    "num_samples_per_task = 10\n",
    "pbar = tqdm(total=len(data) * num_samples_per_task)\n",
    "out_path = {'base': './HumanEval_13b-code_0shot_sample10_bayesian_maxtoken512_maxstep12_base-k0-p0.95-t0.6_rth1.00_double_check2.jsonl',\n",
    "            'sss':  './HumanEval_13b-code_0shot_sample10_bayesian_maxtoken512_maxstep12_sss-k0-p0.95-t0.6_rth1.00_th0.6_double_check2.jsonl',\n",
    "            'sss_autoth1':  './HumanEval_13b-code_0shot_sample10_bayesian_maxtoken512_maxstep12_sss-k0-p0.95-t0.6_rth1.00_autoth0.6-st1-mat0.85-var1e-2--mom1-0.5-mom2-0.9_double_check2.jsonl',\n",
    "            'sss_autoth2':  './HumanEval_13b-code_0shot_sample10_bayesian_maxtoken512_maxstep12_sss-k0-p0.95-t0.6_rth1.00_autoth0.6-st1-mat0.90-var1e-2--mom1-0.5-mom2-0.9_double_check2.jsonl',\n",
    "            }\n",
    "samples = {'base': [], 'sss': [], 'sss_autoth1': [], 'sss_autoth2': []}\n",
    "with open('HumanEval_13b-code_0shot_sample10_bayesian_maxtoken512_maxstep12_sss-k0-p0.95-t0.6_rth1.00_autoth0.6-st1-mat0.85-0.90-var1e-2-mom1-0.5-mom2-0.9_double_check2.txt', 'w') as f:\n",
    "    for i,task_id in enumerate(data):\n",
    "        input_ids = clip_input(tokenizer, data[task_id], task_name)\n",
    "        f.write(f'prompt length: {input_ids.shape[1]} \\n')\n",
    "        for n in range(num_samples_per_task):\n",
    "            if i == 0 and n == 0:\n",
    "                th_stop_draft_sss = 0.6\n",
    "                th_stop_draft_sss_autoth1  = 0.6\n",
    "                th_stop_draft_sss_autoth2  = 0.6\n",
    "            else:\n",
    "                th_stop_draft_sss = result_sss['th_stop_draft']\n",
    "                th_stop_draft_sss_autoth1 = result_sss_autoth1['th_stop_draft']\n",
    "                th_stop_draft_sss_autoth2 = result_sss_autoth2['th_stop_draft']\n",
    "            f.write('sss th: {:.4f}, sss autoth1: {:.4f}, sss autoth2: {:.4f} \\n'.format(th_stop_draft_sss, th_stop_draft_sss_autoth1, th_stop_draft_sss_autoth2))\n",
    "            result_base = infer_input_ids(model, tokenizer, input_ids, generate_fn='base',seed=n, \n",
    "                        max_new_tokens=512, do_sample=True, early_stop=True, \n",
    "                        top_k=0, top_p=0.95, temperature=0.6)\n",
    "            result_sss = infer_input_ids(model, tokenizer, input_ids, generate_fn='sss', seed=n,\n",
    "                        max_new_tokens=512, early_stop=True, max_step_draft=12,\n",
    "                        th_stop_draft=th_stop_draft_sss, th_random_draft=1.0, auto_th_stop_draft=False,\n",
    "                        do_sample=True, do_sample_draft=True, top_k=0, top_p=0.95, temperature=0.6)\n",
    "            result_sss_autoth1 = infer_input_ids(model, tokenizer, input_ids, generate_fn='sss', seed=n,\n",
    "                        max_new_tokens=512, early_stop=True,max_step_draft=12,\n",
    "                        th_stop_draft=th_stop_draft_sss_autoth1, th_random_draft=1.0, auto_th_stop_draft=True, auto_parameters=[1,0.5,0.85,1e-2,0.9],\n",
    "                        do_sample=True, do_sample_draft=True, top_k=0, top_p=0.95, temperature=0.6)\n",
    "            result_sss_autoth2 = infer_input_ids(model, tokenizer, input_ids, generate_fn='sss', seed=n,\n",
    "                        max_new_tokens=512, early_stop=True,max_step_draft=12,\n",
    "                        th_stop_draft=th_stop_draft_sss_autoth2, th_random_draft=1.0, auto_th_stop_draft=True, auto_parameters=[1,0.5,0.90,1e-2,0.9],\n",
    "                        do_sample=True, do_sample_draft=True, top_k=0, top_p=0.95, temperature=0.6)\n",
    "            \n",
    "            results = [\n",
    "                ('base', result_base),\n",
    "                ('sss', result_sss),\n",
    "                ('sss_autoth1', result_sss_autoth1),\n",
    "                ('sss_autoth2', result_sss_autoth2),\n",
    "            ]\n",
    "\n",
    "            for key, result in results:\n",
    "                main_metrics['time_' + key].append(result['time'])\n",
    "                main_metrics['token_time_' + key].append(result['time'] / result['generate_ids'].shape[1])\n",
    "                \n",
    "                if key != 'base':\n",
    "                    main_metrics['matchness_' + key].append(result['matchness'])\n",
    "                    main_metrics['num_drafted_tokens_' + key].append(result['num_drafted_tokens'])\n",
    "                sample=filter_code(fix_indents(result['completion']))\n",
    "                samples[key]+= [{'task_id': task_id, 'completion': sample}]\n",
    "                # f.write(f'{key}, completion: \\n{sample}\\n')\n",
    "\n",
    "            metric = {\n",
    "                'mean time base tem 0.6':np.mean(main_metrics['time_base']),\n",
    "                'mean time sss tem 0.6':np.mean(main_metrics['time_sss']),\n",
    "                'mean time sss autoth rth 1.00 tem 0.6 mat 0.85':np.mean(main_metrics['time_sss_autoth1']),\n",
    "                'mean time sss autoth rth 1.00 tem 0.6 mat 0.90':np.mean(main_metrics['time_sss_autoth2']),\n",
    "                'E2E mean speed up sss tem 0.6':np.mean(main_metrics['time_base'])/np.mean(main_metrics['time_sss']),\n",
    "                'E2E mean speed up sss autoth rth 1.00 tem 0.6 mat 0.85':np.mean(main_metrics['time_base'])/np.mean(main_metrics['time_sss_autoth1']),\n",
    "                'E2E mean speed up sss autoth rth 1.00 tem 0.6 mat 0.90':np.mean(main_metrics['time_base'])/np.mean(main_metrics['time_sss_autoth2']),\n",
    "                'E2E mean token speed up sss tem 0.6':np.mean(main_metrics['token_time_base'])/np.mean(main_metrics['token_time_sss']),  \n",
    "                'E2E mean token speed up sss autoth rth 1.00 tem 0.6 mat 0.85':np.mean(main_metrics['token_time_base'])/np.mean(main_metrics['token_time_sss_autoth1']), \n",
    "                'E2E mean token speed up sss autoth rth 1.00 tem 0.6 mat 0.90':np.mean(main_metrics['token_time_base'])/np.mean(main_metrics['token_time_sss_autoth2']),      \n",
    "                'mean matchness sss tem 0.6':np.mean(main_metrics['matchness_sss']),\n",
    "                'mean matchness sss autoth rth 1.00 tem 0.6 mat 0.85':np.mean(main_metrics['matchness_sss_autoth1']),\n",
    "                'mean matchness sss autoth rth 1.00 tem 0.6 mat 0.90':np.mean(main_metrics['matchness_sss_autoth2']),\n",
    "                'mean num_drafted_tokens sss tem 0.6':np.mean(main_metrics['num_drafted_tokens_sss']),\n",
    "                'mean num_drafted_tokens sss autoth rth 1.00 tem 0.6 mat 0.85':np.mean(main_metrics['num_drafted_tokens_sss_autoth1']),\n",
    "                'mean num_drafted_tokens sss autoth rth 1.00 tem 0.6 mat 0.90':np.mean(main_metrics['num_drafted_tokens_sss_autoth2']),\n",
    "            }\n",
    "            for key, value in metric.items():\n",
    "                if isinstance(value, float):\n",
    "                    metric[key] = f\"{value:.4f}\"\n",
    "            # task_id = x['task_id']\n",
    "            f.write(f'{task_id},sample{n},{metric} \\n')\n",
    "            f.flush()\n",
    "        pbar.update(num_samples_per_task)\n",
    "    for key in out_path.keys():\n",
    "        write_jsonl(out_path[key], samples[key])  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "main_metrics = {'time_base':[], 'time_sss':[], 'time_sss_autoth1':[], 'time_sss_autoth2':[], 'time_sss_autoth3':[],\n",
    "                'token_time_base':[], 'token_time_sss':[], 'token_time_sss_autoth1':[], 'token_time_sss_autoth2':[], 'token_time_sss_autoth3':[],\n",
    "                'matchness_sss':[],'num_drafted_tokens_sss':[],\n",
    "                'matchness_sss_autoth1':[],'num_drafted_tokens_sss_autoth1':[],\n",
    "                'matchness_sss_autoth2':[],'num_drafted_tokens_sss_autoth2':[],\n",
    "                'matchness_sss_autoth3':[],'num_drafted_tokens_sss_autoth3':[],\n",
    "                }\n",
    "num_samples_per_task = 10\n",
    "pbar = tqdm(total=len(data) * num_samples_per_task)\n",
    "out_path = {'base': './HumanEval_13b-code_0shot_sample10_bayesian_maxtoken512_maxstep12_base-k0-p0.95-t0.6_rth1.00.jsonl',\n",
    "            'sss':  './HumanEval_13b-code_0shot_sample10_bayesian_maxtoken512_maxstep12_sss-k0-p0.95-t0.6_rth1.00_th0.6.jsonl',\n",
    "            'sss_autoth1':  './HumanEval_13b-code_0shot_sample10_bayesian_maxtoken512_maxstep12_sss-k0-p0.95-t0.6_rth0.95_autoth0.6-st1-mat0.85-var1e-2--mom1-0.5-mom2-0.9.jsonl',\n",
    "            'sss_autoth2':  './HumanEval_13b-code_0shot_sample10_bayesian_maxtoken512_maxstep12_sss-k0-p0.95-t0.6_rth1.00_autoth0.6-st1-mat0.85-var1e-2--mom1-0.5-mom2-0.9_v2.jsonl',\n",
    "            'sss_autoth3':  './HumanEval_13b-code_0shot_sample10_bayesian_maxtoken512_maxstep12_sss-k0-p0.95-t0.6_rth1.00_autoth0.6-st1-mat0.90-var1e-2--mom1-0.5-mom2-0.9_v2.jsonl',\n",
    "            }\n",
    "samples = {'base': [], 'sss': [], 'sss_autoth1': [], 'sss_autoth2': [], 'sss_autoth3': []}\n",
    "with open('HumanEval_13b-code_0shot_sample10_bayesian_maxtoken512_maxstep12_sss-k0-p0.95-t0.6_rth0.95-1.00_autoth0.6-st1-mat0.85-0.90-var1e-2-mom1-0.5-mom2-0.9.txt', 'w') as f:\n",
    "    for i,task_id in enumerate(data):\n",
    "        input_ids = clip_input(tokenizer, data[task_id], task_name)\n",
    "        f.write(f'prompt length: {input_ids.shape[1]} \\n')\n",
    "        for n in range(num_samples_per_task):\n",
    "            if i == 0 and n == 0:\n",
    "                th_stop_draft_sss = 0.6\n",
    "                th_stop_draft_sss_autoth1  = 0.6\n",
    "                th_stop_draft_sss_autoth2  = 0.6\n",
    "                th_stop_draft_sss_autoth3  = 0.6\n",
    "            else:\n",
    "                th_stop_draft_sss = result_sss['th_stop_draft']\n",
    "                th_stop_draft_sss_autoth1 = result_sss_autoth1['th_stop_draft']\n",
    "                th_stop_draft_sss_autoth2 = result_sss_autoth2['th_stop_draft']\n",
    "                th_stop_draft_sss_autoth3 = result_sss_autoth3['th_stop_draft']\n",
    "            f.write('sss th: {:.4f}, sss autoth1: {:.4f}, sss autoth2: {:.4f}, sss autoth3: {:.4f} \\n'.format(th_stop_draft_sss, th_stop_draft_sss_autoth1, th_stop_draft_sss_autoth2, th_stop_draft_sss_autoth3))\n",
    "            result_base = infer_input_ids(model, tokenizer, input_ids, generate_fn='base',seed=n, \n",
    "                        max_new_tokens=512, do_sample=True, early_stop=True, \n",
    "                        top_k=0, top_p=0.95, temperature=0.6)\n",
    "            result_sss = infer_input_ids(model, tokenizer, input_ids, generate_fn='sss', seed=n,\n",
    "                        max_new_tokens=512, early_stop=True, max_step_draft=12,\n",
    "                        th_stop_draft=th_stop_draft_sss, th_random_draft=1.0, auto_th_stop_draft=False,\n",
    "                        do_sample=True, do_sample_draft=True, top_k=0, top_p=0.95, temperature=0.6)\n",
    "            result_sss_autoth1 = infer_input_ids(model, tokenizer, input_ids, generate_fn='sss', seed=n,\n",
    "                        max_new_tokens=512, early_stop=True,max_step_draft=12,\n",
    "                        th_stop_draft=th_stop_draft_sss_autoth1, th_random_draft=0.95, auto_th_stop_draft=True, auto_parameters=[1,0.5,0.85,1e-2,0.9],\n",
    "                        do_sample=True, do_sample_draft=True, top_k=0, top_p=0.95, temperature=0.6)\n",
    "            result_sss_autoth2 = infer_input_ids(model, tokenizer, input_ids, generate_fn='sss', seed=n,\n",
    "                        max_new_tokens=512, early_stop=True,max_step_draft=12,\n",
    "                        th_stop_draft=th_stop_draft_sss_autoth2, th_random_draft=1.0, auto_th_stop_draft=True, auto_parameters=[1,0.5,0.85,1e-2,0.9],\n",
    "                        do_sample=True, do_sample_draft=True, top_k=0, top_p=0.95, temperature=0.6)\n",
    "            result_sss_autoth3 = infer_input_ids(model, tokenizer, input_ids, generate_fn='sss', seed=n,\n",
    "                        max_new_tokens=512, early_stop=True,max_step_draft=12,\n",
    "                        th_stop_draft=th_stop_draft_sss_autoth3, th_random_draft=1.0, auto_th_stop_draft=True, auto_parameters=[1,0.5,0.90,1e-2,0.9],\n",
    "                        do_sample=True, do_sample_draft=True, top_k=0, top_p=0.95, temperature=0.6)\n",
    "            \n",
    "            results = [\n",
    "                ('base', result_base),\n",
    "                ('sss', result_sss),\n",
    "                ('sss_autoth1', result_sss_autoth1),\n",
    "                ('sss_autoth2', result_sss_autoth2),\n",
    "                ('sss_autoth3', result_sss_autoth3),\n",
    "            ]\n",
    "\n",
    "            for key, result in results:\n",
    "                main_metrics['time_' + key].append(result['time'])\n",
    "                main_metrics['token_time_' + key].append(result['time'] / result['generate_ids'].shape[1])\n",
    "                \n",
    "                if key != 'base':\n",
    "                    main_metrics['matchness_' + key].append(result['matchness'])\n",
    "                    main_metrics['num_drafted_tokens_' + key].append(result['num_drafted_tokens'])\n",
    "                sample=filter_code(fix_indents(result['completion']))\n",
    "                samples[key]+= [{'task_id': task_id, 'completion': sample}]\n",
    "                # f.write(f'{key}, completion: \\n{sample}\\n')\n",
    "\n",
    "            metric = {\n",
    "                'mean time base tem 0.6':np.mean(main_metrics['time_base']),\n",
    "                'mean time sss tem 0.6':np.mean(main_metrics['time_sss']),\n",
    "                'mean time sss autoth rth 0.95 tem 0.6 mat 0.85':np.mean(main_metrics['time_sss_autoth1']),\n",
    "                'mean time sss autoth rth 1.00 tem 0.6 mat 0.85':np.mean(main_metrics['time_sss_autoth2']),\n",
    "                'mean time sss autoth rth 1.00 tem 0.6 mat 0.90':np.mean(main_metrics['time_sss_autoth3']),\n",
    "                'E2E mean speed up sss tem 0.6':np.mean(main_metrics['time_base'])/np.mean(main_metrics['time_sss']),\n",
    "                'E2E mean speed up sss autoth rth 0.95 tem 0.6 mat 0.85':np.mean(main_metrics['time_base'])/np.mean(main_metrics['time_sss_autoth1']),\n",
    "                'E2E mean speed up sss autoth rth 1.00 tem 0.6 mat 0.85':np.mean(main_metrics['time_base'])/np.mean(main_metrics['time_sss_autoth2']),\n",
    "                'E2E mean speed up sss autoth rth 1.00 tem 0.6 mat 0.90':np.mean(main_metrics['time_base'])/np.mean(main_metrics['time_sss_autoth3']),\n",
    "                'E2E mean token speed up sss tem 0.6':np.mean(main_metrics['token_time_base'])/np.mean(main_metrics['token_time_sss']),\n",
    "                'E2E mean token speed up sss autoth rth 0.95 tem 0.6 mat 0.85':np.mean(main_metrics['token_time_base'])/np.mean(main_metrics['token_time_sss_autoth1']),   \n",
    "                'E2E mean token speed up sss autoth rth 1.00 tem 0.6 mat 0.85':np.mean(main_metrics['token_time_base'])/np.mean(main_metrics['token_time_sss_autoth2']), \n",
    "                'E2E mean token speed up sss autoth rth 1.00 tem 0.6 mat 0.90':np.mean(main_metrics['token_time_base'])/np.mean(main_metrics['token_time_sss_autoth3']),      \n",
    "                'mean matchness sss tem 0.6':np.mean(main_metrics['matchness_sss']),\n",
    "                'mean matchness sss autoth rth 0.95 tem 0.6 mat 0.85':np.mean(main_metrics['matchness_sss_autoth1']),\n",
    "                'mean matchness sss autoth rth 1.00 tem 0.6 mat 0.85':np.mean(main_metrics['matchness_sss_autoth2']),\n",
    "                'mean matchness sss autoth rth 1.00 tem 0.6 mat 0.90':np.mean(main_metrics['matchness_sss_autoth3']),\n",
    "                'mean num_drafted_tokens sss tem 0.6':np.mean(main_metrics['num_drafted_tokens_sss']),\n",
    "                'mean num_drafted_tokens sss autoth rth 0.95 tem 0.6 mat 0.85':np.mean(main_metrics['num_drafted_tokens_sss_autoth1']),\n",
    "                'mean num_drafted_tokens sss autoth rth 1.00 tem 0.6 mat 0.85':np.mean(main_metrics['num_drafted_tokens_sss_autoth2']),\n",
    "                'mean num_drafted_tokens sss autoth rth 1.00 tem 0.6 mat 0.90':np.mean(main_metrics['num_drafted_tokens_sss_autoth3']),\n",
    "            }\n",
    "            for key, value in metric.items():\n",
    "                if isinstance(value, float):\n",
    "                    metric[key] = f\"{value:.4f}\"\n",
    "            # task_id = x['task_id']\n",
    "            f.write(f'{task_id},sample{n},{metric} \\n')\n",
    "            f.flush()\n",
    "        pbar.update(num_samples_per_task)\n",
    "    for key in out_path.keys():\n",
    "        write_jsonl(out_path[key], samples[key])  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "results_name = './HumanEval_13b-code_0shot_sample10_bayesian_maxtoken512_maxstep12_base-k0-p0.95-t0.6_rth1.00_double_check.jsonl'\n",
    "results_name2 = './HumanEval_13b-code_0shot_sample10_bayesian_maxtoken512_maxstep12_base-k0-p0.95-t0.6_rth1.00_double_check_output.jsonl'\n",
    "\n",
    "input_filename = results_name\n",
    "output_filename = results_name2\n",
    "\n",
    "with open(input_filename, 'r') as input_file, open(output_filename, 'w') as output_file:\n",
    "    for line in input_file:\n",
    "        try:\n",
    "            data = json.loads(line)\n",
    "            completion = data.get(\"completion\", \"\")\n",
    "            data[\"completion\"] = \"    \"+completion\n",
    "            output_file.write(json.dumps(data) + '\\n')\n",
    "        except json.JSONDecodeError:\n",
    "            print(f\"Error decoding JSON: {line}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!git clone https://github.com/openai/human-eval\n",
    "!pip install -e human-eval"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n",
    "!evaluate_functional_correctness ./HumanEval_13b-code_0shot_sample10_bayesian_maxtoken512_maxstep12_base-k0-p0.95-t0.6_rth1.00_double_check_output.jsonl"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "skipgpt",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.9"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
